# -*- coding: utf-8 -*-
"""
Created on Tue Nov 23 10:51:57 2021

@author: zhiyinan
"""
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import Parameters_1D as par_RNN 
from Polar1D_vectorized import *
import van_Genuchten1D_vectorized as vg 
from scipy import integrate
from sklearn.preprocessing import MinMaxScaler
from casadi import *
import matplotlib.font_manager as font_manager
import time
import pandas as pd
from sklearn.linear_model import Ridge
import random

# %% Define NN layers
def sigmoid(x):
    return 1/(1+np.exp(-x))

def LSTMlayer_sig(warr, uarr, barr, x_t,h_tm1,c_tm1):
    
    '''
    c_tm1 = np.array([0,0]).reshape(1,2)
    h_tm1 = np.array([0,0]).reshape(1,2)
    x_t   = np.array([1]).reshape(1,1)
    
    warr.shape = (nfeature,hunits*4)
    uarr.shape = (hunits,hunits*4)
    barr.shape = (hunits*4,)
    
    '''
    
    s_t = (mtimes(x_t,warr) + mtimes(h_tm1,uarr) + barr.reshape((1, barr.shape[0])))
    hunit = uarr.shape[0]
    i  = sigmoid(s_t[:,:hunit])
    f  = sigmoid(s_t[:,1*hunit:2*hunit])
    # _c = np.tanh(s_t[:,2*hunit:3*hunit])
    _c = sigmoid(s_t[:,2*hunit:3*hunit])
    o  = sigmoid(s_t[:,3*hunit:])
    c_t = i*_c + f*c_tm1
    # h_t = o*np.tanh(c_t)
    h_t = o*sigmoid(c_t)
    return h_t, c_t

def LSTMlayer_tanh(warr, uarr, barr, x_t,h_tm1,c_tm1):
    
    '''
    c_tm1 = np.array([0,0]).reshape(1,2)
    h_tm1 = np.array([0,0]).reshape(1,2)
    x_t   = np.array([1]).reshape(1,1)
    
    warr.shape = (nfeature,hunits*4)
    uarr.shape = (hunits,hunits*4)
    barr.shape = (hunits*4,)
    
    '''
    
    s_t = (mtimes(x_t,warr) + mtimes(h_tm1,uarr) + barr.reshape((1, barr.shape[0])))
    hunit = uarr.shape[0]
    i  = sigmoid(s_t[:,:hunit])
    f  = sigmoid(s_t[:,1*hunit:2*hunit])
    _c = np.tanh(s_t[:,2*hunit:3*hunit])
    # _c = sigmoid(s_t[:,2*hunit:3*hunit])
    o  = sigmoid(s_t[:,3*hunit:])
    c_t = i*_c + f*c_tm1
    h_t = o*np.tanh(c_t)
    # h_t = o*sigmoid(c_t)
    return h_t, c_t




# %% Define Model 1
model_1 = tf.keras.models.load_model("Modeling/p20_f1_10s_sig_2h_012_027_5_150_e8_n")
#  "Modeling/p20_f1_10s_sig_2h_012_027_5_150_e8"
def M1(x):
    weightLSTM_1 = model_1.layers[0].get_weights()
    warr1, uarr1, barr1 = weightLSTM_1

    weightD = model_1.layers[1].get_weights()
    kernel2, bias2 = weightD
    
    n1 = 10
    # Initialise the model with the hidden states
    h_t1 = SX.zeros((1, n1))
    # Long term memory
    C_t1 = SX.zeros((1, n1))
    for i in range(x.shape[0]):
        xi = x[i,:].reshape((1,2))
        L1 = LSTMlayer_sig(warr1, uarr1, barr1, xi, h_t1[i,:], C_t1[i,:])
        h_t1 = vertcat(h_t1, L1[0])
        C_t1 = vertcat(C_t1, L1[1])
    
    #  array([-2.7764497, -2.6578658, -2.6829786, -2.7356846, -2.73696  ]
    # Dense layer
    out2 = []
    for i in range(bias2.shape[0]):
        temp2 = mtimes(h_t1[-1,:], kernel2)[i]
    temp2 = bias2.reshape((1, bias2.shape[0])) + mtimes(h_t1[-1,:], kernel2)
    out2 = sigmoid(temp2)     
    
    return out2

# %% Input Signal generation
np.random.seed(5)
random.seed(5)


soilPars = loamySoil() # The soil parameters
totalDepth, axialNodes, axialDistance, totalNodes, numOfActuators, interPoints = par_RNN.spatialVariables_1D()  #Spatial parameters
samplingTime, samplingTimeInternal, internalTimeSteps, totTimeSteps = \
                                 par_RNN.temporalVariable_dataGen(1000*2*60*60) #time parameters

# lbu = -5.10e-05  #  unit [m/s]
# ubu = -4.94e-05  #  unit [m/s]

# irr_freq = int(20*60*60/samplingTime)   # PRBS change once only if a minimum of 100 hr is reached

u = np.zeros((totTimeSteps,1))


i=0
irr_freq = int(20*60*60/samplingTime)
while i <= 500:
    ulist = np.array([-0.5e-6, -0.8e-7, -0.5e-7, -0.7e-7, -0.1e-6, -0.2e-6, -0.3e-6, -0.6e-6, -0.4e-6, -1e-6, -1.5e-6])/5
    # ulist = [-2/1000/3600, -1.5/1000/3600, -0.9/1000/3600, -1.2/1000/3600, -0.5/1000/3600]
    t = irr_freq*(1 + random.randrange(3)) + random.randrange(irr_freq)
    u[i:i+t] = random.choice(ulist)
    i = i+t+int(random.randrange(irr_freq*1.8)*0.8)
    print(i)
else:  
    while i<=1000:
        ulist = np.array([-0.5e-6, -0.8e-7, -0.5e-7, -0.7e-7, -0.1e-6, -0.2e-6, -0.3e-6, -0.6e-6, -0.4e-6, -1e-6, -1.5e-6])/5
        # ulist = [-2/1000/3600, -1.5/1000/3600, -0.9/1000/3600, -1.2/1000/3600, -0.5/1000/3600]
        t = irr_freq*(1 + random.randrange(3)) + random.randrange(irr_freq)
        u[i:i+t] = random.choice(ulist)
        i = i+t
        print(i)
    

plt.figure(1)
plt.plot(u)                


initCondition = -0.9*np.ones(totalNodes)
#Arrays to store the results




def ODE(t, head, irrigAmount, cropCoeff, refEvap, soilPars):
    return Richards1D_vectorized(head, irrigAmount, cropCoeff, refEvap, soilPars)

cropCoeff = 0.88*np.ones(totTimeSteps)
refEvap = 3.102e-08*np.ones(totTimeSteps)

  # %% Integration 

headArray= np.zeros((totTimeSteps+1, totalNodes)) # Pressure head, the state
headArray[0,:] = initCondition

volMoisture = np.zeros((totTimeSteps+1, totalNodes)) # Volumetric moisture, the measured output
volMoisture[0,:] = volMoistureAllNodes_1D(initCondition, soilPars).ravel()


i=1
temp = headArray[i-1]
ui = u[i-1]
for t in range(int(samplingTime/samplingTimeInternal)): 
    sol = integrate.solve_ivp(ODE, args = (ui, cropCoeff[i-1], refEvap[i-1], soilPars),
                              t_span=[0,samplingTimeInternal],y0=tuple(temp),method='BDF')
    temp = sol.y[:,-1]
    # ui = 0
headArray[i,:]=sol.y[:,-1]
volMoisture[i,:] = volMoistureAllNodes_1D(headArray[i,:], soilPars).ravel()

for i in range(2, totTimeSteps + 1):  # totTimeSteps + 1
    print(i)
    temp = headArray[i-1]
    ui = u[i-1]
    for t in range(int(samplingTime/samplingTimeInternal)): 
        sol = integrate.solve_ivp(ODE, args = (ui, cropCoeff[i-1], refEvap[i-1], soilPars),
                              t_span=[0,samplingTimeInternal],y0=tuple(temp),method='BDF')
        temp = sol.y[:,-1]
    # print(sol.y[:,-1])
        # ui = 0
    headArray[i,:]=sol.y[:,-1]
    volMoisture[i,:] = volMoistureAllNodes_1D(headArray[i,:], soilPars).ravel()

plt.figure(2)
plt.plot(volMoisture[:1000,19])

# %%

def multivariate_data(dataset, target, start_index, end_index, history_size,
                      target_size,u_shape): #, single_step=False):
    data = []
    labels = []

    start_index = start_index + history_size
    if end_index is None:
        end_index = len(dataset) - target_size
        
    for i in range(start_index, end_index):
        temp_p = dataset[i-history_size:i,:]
        for j in range(target_size-1):
        # temp_p = np.concatenate(((dataset[i-history_size:i,:]).flatten(), 
        #                            (dataset[i:i+target_size-1,:u_shape]).flatten()))
            temp = np.array([dataset[i+j,0], dataset[i,1]]).reshape((1,2))
            temp_p = np.concatenate((temp_p,temp))
        data.append(temp_p)
        labels.append((target[i:i+target_size].flatten()))
    

    return np.array(data), np.array(labels)


# %% Training data
past = 20
future = 1
y_index = 19


y = volMoisture[:29001,y_index]
yn_train = 0.1*(max(y) - min(y))*np.random.choice([-1,1],(y.shape))*np.random.random((y.shape))
y_n = y*(1+yn_train)

y_train = y[1:]
x_train = np.zeros((29000,2))
x_train[:,0] = u[:29000,0]
x_train[:,1] = y_n[:-1]

scaler_x = MinMaxScaler()
x_train = scaler_x.fit_transform(x_train)
scaler_y = MinMaxScaler()
y_train = scaler_y.fit_transform(y_train.reshape((y_train.shape[0],1)))

x_data, y_data = multivariate_data(x_train, y_train, 0, None, past, future,1)

OUTPUT_SIZE = 1
# %% Identifying the Model
model_1 = tf.keras.Sequential() 
model_1.add(tf.keras.layers.LSTM(10, input_shape=(past+future-1,2), activation = "sigmoid"))
# model_1.add(tf.keras.layers.LSTM(3, activation = "sigmoid", return_sequences=True))
# model_1.add(tf.keras.layers.LSTM(5, activation = "sigmoid"))
# model_1.add(tf.keras.layers.Dense(5, input_shape=(past*2+future-1,), activation = "sigmoid"))  #, activation = 'sigmoid'
# model.add(Dense(50))
# model.add(tf.keras.layers.SimpleRNN(10, activation = 'tanh', return_sequences=True))
# model_1.add(tf.keras.layers.Dense(20, activation = 'sigmoid')) 
# model_1.add(tf.keras.layers.Dense(20, activation = 'sigmoid')) 
# model_1.add(tf.keras.layers.LSTM(20, return_sequences=True)) 
# model_1.add(tf.keras.layers.LSTM(5, activation = "sigmoid"))  #, return_sequences=True
# model_1.add(tf.keras.layers.Dense(20, activation = 'tanh'))  # , activation = 'sigmoid',activation="selu"
# model.add(tf.keras.layers.Dense(1, activation = 'selu'))
# add output layer
# model_2.add(tf.keras.layers.Dense(5, activation = 'tanh')) 
model_1.add(tf.keras.layers.Dense(OUTPUT_SIZE, activation = "sigmoid"))

opt = tf.keras.optimizers.Adam(learning_rate=0.005)
model_1.compile(optimizer=opt, loss='mse', metrics=['mean_squared_error'])

epoch = 500

callback = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=3)
Dense_history = model_1.fit(x_data,y_data.reshape((y_data.shape[0],y_data.shape[1])), epochs=epoch,
                               batch_size=200,
                                validation_split=0.15,
                           # steps_per_epoch = steps_per_epoch,
                                # validation_data = (x_t, y_t),
                                # validation_steps = validation_steps,
                               callbacks = [callback])


# %% Validation data
# uv = u.flatten()
uv = np.loadtxt("Modeling/u_vali_20hrIF_1000_013_027.txt")
volMoisture = np.loadtxt("Modeling/y_vali_20hrIF_1000_013_027.txt")
yv = volMoisture[:,y_index]
# uv = u[29000:]
# yv = volMoisture[29000:,y_index]
a = 900
b = 1000
yn = 0.2*(max(yv) - min(yv))*np.random.choice([-1,1],(yv.shape))*np.random.random((yv.shape))
yv_n = yv*(1+yn)
y_test = yv[1:b+1]
x_test = np.zeros((b,2))
x_test[:,0] = uv[:b]
x_test[:,1] = yv_n[:b]

scaler_xt = MinMaxScaler()
x_test = scaler_xt.fit_transform(x_test)
y_test = y_test.reshape((y_test.shape[0],1))
scaler_yt = MinMaxScaler()
y_test = scaler_yt.fit_transform(y_test)



x_t, y_t = multivariate_data(x_test, y_test, 0, None, past, future,1)

# %%
y_pt_1 = model_1.predict(x_t)
ypt_1 = scaler_yt.inverse_transform(y_pt_1)
y_act = scaler_yt.inverse_transform(y_t.reshape((y_t.shape[0], y_t.shape[1])))

csfont = {'fontname':'Times New Roman'}
font = font_manager.FontProperties(family='Times New Roman',
                                    # weight='bold',
                                    style='normal', size=15)
plt.figure(3)
plt.plot(y_act[:,0], label = 'Actual Data')
plt.plot(ypt_1[:,0], label = 'Single-step-ahead Prediction')
# plt.plot(ypt_1[-500:,-1], label = 'Multi-step-ahead NN')
# plt.plot(ypt_2[30:,0], label = 'Multi-step-ahead NN -1st')
# plt.plot(ypt_2[:,-1], label = 'Multi-step-ahead NN - 30th')
plt.xlabel('Time Steps', **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)
plt.legend(frameon=False, ncol=1, prop=font)

# %% Multi-step ahead prediction performance
a = 900
t_pred = 20  # 30 sampling times 
# umax = max(uv[:])
# umin = min(uv[:])
# u_sc = (uv[:,0] - umin)/(umax - umin)
error = np.zeros(x_t.shape[0])
y_save = np.zeros((a,t_pred))

for t0 in range(a):
    print(t0)
    y_sim_1 = np.zeros((t_pred,1))
    NN0_1 = x_t[t0,:]
    # y_sub = np.array([DM(M1(NN0_1)), DM(M2(NN0_1)), DM(M1(NN0_1))])  #, DM(M4(NN0_1)), DM(M5(NN0_1))])
    y_sim_1[0,:] = np.array(DM(M1(NN0_1))).flatten()
    u_sim = x_t[t0:t0+t_pred,-1,0]
    for i in range(t_pred-1):
        NN0_1 = np.concatenate((NN0_1[1:], np.array([u_sim[i], y_sim_1[i,0]]).reshape((1,2))))
        # y_sub = np.array([DM(M1(NN0_1)), DM(M2(NN0_1)), DM(M1(NN0_1))])  #, DM(M4(NN0_1)), DM(M5(NN0_1))])
        y_sim_1[i+1,:] = np.array(DM(M1(NN0_1))).flatten()
    yp_1 = scaler_yt.inverse_transform(y_sim_1).flatten()
    y_save[t0,:] = yp_1
    error[t0] = sum((y_act[t0:t0+t_pred].flatten() - yp_1)**2/y_act[t0:t0+t_pred].flatten())
# yp_2 = scaler_yt.inverse_transform(y_sim_2)
# yp = (yp_1 + yp_2)/2
np.savetxt("Modeling/y_pred_multi_20hrIF_1000_013_027.txt", y_save)
# %%
y_save = np.loadtxt("Modeling/y_pred_multi_20hrIF_1000_013_027.txt")
font = font_manager.FontProperties(family='Times New Roman',
                                   # weight='bold',
                                    style='normal', size=17)
tplot = np.arange(a)
b = 330
c = 390

fig,ax1 = plt.subplots()
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
ax1.plot(tplot, y_act[:a], label = 'Actual Trajectory', color = "k", linewidth = 3)
ax1.plot(tplot, y_save[:,0], label = "Single-step-ahead") # 0.000914381408118862
# ax1.plot(tplot[:], y_save[:,4], label = "5-step-ahead", color = "r")  # 0.0017256467027993415
ax1.plot(tplot[9:], y_save[:a-9,9], label = "Ten-step-ahead", color = "r", linestyle = "dashdot") # 0.002444138910936656
ax1.plot(tplot[19:], y_save[:a-19,19], label = "Twenty-step-ahead", color = "y", linestyle = "dashed") # 0.0030556959486465352
ax1.legend(frameon=False, bbox_to_anchor=(0.2,0.7,0.3,0.3), ncol=1, prop=font)
plt.xlabel("Time Steps", **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)


# ax2 = plt.axes([0,0,1,1])
# ip = InsetPosition(ax1, [0.15,0.75,0.2,0.2])
# ax2.set_axes_locator(ip)
# mark_inset(ax1, ax2, loc1=2, loc2=4, fc="none", ec='0.6')
# ax2.plot(tplot[b:c+20], y_act[b:c+20], label = 'Actual Trajectory', color = "k", linewidth = 3)
# ax2.plot(tplot[b:c], y_save[b:c,0], label = "Single-step-ahead") 
# ax2.plot(tplot[b+9:c+9], y_save[b:c,9], label = "Ten-step-ahead", color = "r", linestyle = "dashdot") 
# ax2.plot(tplot[b+19:c+19], y_save[b:c,19], label = "Twenty-step-ahead", color = "y", linestyle = "dashed") 
# # ax2.legend(loc=0)
# # ax2.set_xticklabels(ax2.get_xticks(), backgroundcolor='w')
# ax2.tick_params(axis='x', which='major', pad=8)
# plt.suptitle('Composite Plot with original subsystems(Ca)', fontsize=15)
plt.show()

e_0 = np.sqrt(sum((y_act[:900,0] - y_save[:,0])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0])) #(max((y_act[:900,:] - y_save[:,:]).flatten()))
e_5 = np.sqrt(sum((y_act[4:904,0] - y_save[:,4])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0])) 
e_10 = np.sqrt(sum((y_act[9:909,0] - y_save[:,9])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0]))
e_20 = np.sqrt(sum((y_act[19:919,0] - y_save[:,19])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0]))

print(e_0*100)
print(e_5*100)
print(e_10*100)
print(e_20*100)
# e_0 = np.sqrt(sum((y_act[:900,0] - y_save[:,0])**2/y_act[:900,0]**2))/900
# e_5 = np.sqrt(sum((y_act[4:904,0] - y_save[:,4])**2/y_act[4:904,0]**2))/900
# e_10 = np.sqrt(sum((y_act[9:909,0] - y_save[:,9])**2/y_act[9:909,0]**2))/900
# e_20 = np.sqrt(sum((y_act[19:919,0] - y_save[:,19])**2/y_act[19:919,0]**2))/900

# np.savetxt("Modeling/u_vali_20hrIF_1000_013_027.txt", uv)
# np.savetxt("x_vali_20hrIF_1000_013_027.txt", headArray)
# np.savetxt("Modeling/y_vali_20hrIF_1000_013_027.txt", volMoisture)
# np.savetxt("y_pred_multi_20hrIF_1000_013_027.txt", y_save)
